#ifdef __SSE2__
#include <emmintrin.h>
#endif

#ifdef __ARM_NEON__
#include <arm_neon.h>
#endif

#include "vio/front_end/feature_alignment.h"

namespace vio {
namespace feature_alignment {

#define SUBPIX_VERBOSE 0

bool Align1D(const cv::Mat& cur_img,
             const Eigen::Vector2f&
                 dir,  // direction in which the the patch is allowed to move
             uint8_t* ref_patch_with_border,
             uint8_t* ref_patch,
             const int n_iter,
             Eigen::Vector2d& cur_px_estimate,
             double& h_inv) {
  // patch params
  const int half_patch_size = 4;
  const int patch_size = 8;
  const int patch_area = 64;
  bool converged = false;

  // compute drivative of template and prepare inverse compositional
  float __attribute__((__aligned__(16))) ref_patch_dv[patch_area];
  Eigen::Matrix2f H;
  H.setZero();

  // compute gradient and hessian
  const int ref_step = patch_size + 2;  // patch with border
  float* it_dv = ref_patch_dv;
  for (int y = 0; y < patch_size; ++y) {
    // skip border part in patch with border
    uint8_t* it = ref_patch_with_border + (y + 1) * ref_step + 1;
    for (int x = 0; x < patch_size; ++x, ++it, ++it_dv) {
      Eigen::Vector2f J;
      // adjacent pixels: left, right, up, down
      J[0] = 0.5 * (dir[0] * (it[1] - it[-1]) +
                    dir[1] * (it[ref_step] - it[-ref_step]));
      J[1] = 1;
      *it_dv = J[0];
      H += J * J.transpose();
    }
  }
  h_inv = 1.0 / H(0, 0) * patch_size * patch_size;
  Eigen::Matrix2f H_inv = H.inverse();
  float mean_diff = 0;

  // compute pixel location in new image
  float u = cur_px_estimate.x();
  float v = cur_px_estimate.y();

  // termination condition
  const float min_update_squred = 0.03 * 0.03;
  const int cur_step = cur_img.step.p[0];
  float chi2 = 0;
  Eigen::Vector2f update;
  update.setZero();
  for (int iter = 0; iter < n_iter; ++iter) {
    int u_r = std::floor(u);
    int v_r = std::floor(v);
    if (u_r < half_patch_size || v_r < half_patch_size ||
        u_r >= cur_img.cols - half_patch_size ||
        v_r >= cur_img.rows - half_patch_size) {
      break;
    }
    // TODO very rarely this can happen , maybe H is sigular?
    // should not be at corner.. check
    if (std::isnan(u) || std::isnan(v)) {
      return false;
    }
    // compute interpolation weights
    float subpix_x = u - u_r;
    float subpix_y = v - v_r;
    float wTL = (1.0 - subpix_x) * (1.0 - subpix_y);  // Top-Left
    float wTR = subpix_x * (1.0 - subpix_y);
    float wBL = (1.0 - subpix_x) * subpix_y;
    float wBR = subpix_x * subpix_y;

    // loop through search_patch, interplate
    uint8_t* it_ref = ref_patch;
    float* it_ref_dv = ref_patch_dv;
    float new_chi2 = 0.0;
    Eigen::Vector2f Jres;
    Jres.setZero();
    for (int y = 0; y < patch_size; ++y) {
      uint8_t* it = (uint8_t*)cur_img.data +
                    (v_r + y - half_patch_size) * cur_step + u_r -
                    half_patch_size;
      for (int x = 0; x < patch_size; ++x, ++it, ++it_ref, ++it_ref_dv) {
        float search_pixel = wTL * it[0] + wTR * it[1] + wBL * it[cur_step] +
                             wBR * it[cur_step + 1];
        float res = search_pixel - *it_ref + mean_diff;
        Jres[0] -= res * (*it_ref_dv);
        Jres[1] -= res;
        new_chi2 += res * res;
      }
    }
    if (iter > 0 && new_chi2 > chi2) {
      LOG(INFO) << "error increased.";
      u -= update[0];
      v -= update[1];
      break;
    }

    chi2 = new_chi2;
    update = H_inv * Jres;
    u += update[0] * dir[0];
    v += update[0] * dir[1];
    // why: estimate two patches's intensity difference.
    mean_diff += update[1];
    LOG(INFO) << "Iter " << iter << ":"
              << "\t u = " << u << ", v = " << v << "\t update = [" << update[0]
              << ", " << update[1] << "]"
              << "\t new chi2 = " << new_chi2;
    if (update[0] * update[0] + update[1] * update[1] < min_update_squred) {
      LOG(INFO) << "converged.";
      converged = true;
      break;
    }
  }
  cur_px_estimate << u, v;
  return converged;
}

bool Align2D(const cv::Mat& cur_img,
             uint8_t* ref_patch_with_border,
             uint8_t* ref_patch,
             const int n_iter,
             Eigen::Vector2d& cur_px_estimate,
             bool no_simd) {
  const int half_patch_size = 4;
  const int patch_size = 8;
  const int patch_area = 64;
  bool converged = false;

  // compute derivative of template and prepare inverse compositional
  float __attribute__((__aligned__(16))) ref_patch_dx[patch_area];
  float __attribute__((__aligned__(16))) ref_patch_dy[patch_area];
  Eigen::Matrix3f H;
  H.setZero();

  // compute gradient and hessian
  const int ref_step = patch_size + 2;
  float* it_dx = ref_patch_dx;
  float* it_dy = ref_patch_dy;
  for (int y = 0; y < patch_size; y++) {
    uint8_t* it = ref_patch_with_border + (y + 1) * ref_step + 1;
    for (int x = 0; x < patch_size; ++x, ++it, ++it_dx, ++it_dy) {
      Eigen::Vector3f J;
      J[0] = 0.5 * (it[1] - it[-1]);
      J[1] = 0.5 * (it[ref_step] - it[-ref_step]);
      J[2] = 1.0;
      *it_dx = J[0];
      *it_dy = J[1];
      H += J * J.transpose();
    }
  }
  Eigen::Matrix3f H_inv = H.inverse();
  float mean_diff = 0;  // estimates intensity diff mean

  // compute pixel location in new image
  float u = cur_px_estimate.x();
  float v = cur_px_estimate.y();

  // termination condition
  const float min_update_squared = 0.03 * 0.03;
  const int cur_step = cur_img.step.p[0];
  float chi2 = 0;
  Eigen::Vector3f update;
  update.setZero();
  for (int iter = 0; iter < n_iter; iter++) {
    int u_r = std::floor(u);
    int v_r = std::floor(v);
    if (u_r < half_patch_size || v_r < half_patch_size ||
        u_r >= cur_img.cols - half_patch_size ||
        v_r >= cur_img.rows - half_patch_size) {
      break;
    }
    // TODO very rarely this can happen, maybe H is
    // singular? should not be at corner.. check
    if (std::isnan(u) || std::isnan(v)) {
      return false;
    }

    // compute interpolation weights
    float subpix_x = u - u_r;
    float subpix_y = v - v_r;
    float wTL = (1.0 - subpix_x) * (1.0 - subpix_y);
    float wTR = subpix_x * (1.0 - subpix_y);
    float wBL = (1.0 - subpix_x) * subpix_y;
    float wBR = subpix_x * subpix_y;

    // loop through search_patch, interpolate
    uint8_t* it_ref = ref_patch;
    float* it_ref_dx = ref_patch_dx;
    float* it_ref_dy = ref_patch_dy;
    float new_chi2 = 0.0;

    Eigen::Vector3f Jres;
    Jres.setZero();
    for (int y = 0; y < patch_size; y++) {
      uint8_t* it = (uint8_t*)cur_img.data +
                    (v_r + y - half_patch_size) * cur_step + u_r -
                    half_patch_size;
      for (int x = 0; x < patch_size;
           ++x, ++it, ++it_ref, ++it_ref_dx, ++it_ref_dy) {
        float search_pixel = wTL * it[0] + wTR * it[1] + wBL * it[cur_step] +
                             wBR * it[cur_step + 1];
        float res = search_pixel - *it_ref + mean_diff;
        Jres[0] -= res * (*it_ref_dx);
        Jres[1] -= res * (*it_ref_dy);
        Jres[2] -= res;
        new_chi2 += res * res;
      }
    }
    if (iter > 0 && new_chi2 > chi2) {
      LOG(INFO) << "error increased.";
      u -= update[0];
      v -= update[1];
      break;
    }

    update = H_inv * Jres;
    u += update[0];
    v += update[1];
    mean_diff += update[2];
    LOG(INFO) << "Iter " << iter << ":"
              << "\t u = " << u << ", v = " << v << "\t update = [" << update[0]
              << ", " << update[1] << ", " << update[2] << "]"
              << "\t new chi2 = " << new_chi2;

    if (update[0] * update[0] + update[1] * update[1] < min_update_squared) {
      LOG(INFO) << "converged.";
      converged = true;
      break;
    }
  }
  cur_px_estimate << u, v;
  return converged;
}
}  // namespace feature_alignment
}  // namespace vio
